package com.study.bigdata.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Encoder, Encoders, Row, SparkSession, functions}

object Spark03_SparkSQL_UDAF_1 {
  def main(args: Array[String]): Unit = {
    // TODO 创建SparkSql的运行环境   环境+对象+隐式转换
    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("sparkSQL")
    val spark = SparkSession.builder().config(sparkConf).getOrCreate()
    // TODO 执行逻辑
    val df: DataFrame = spark.read.json("data/user.json")
    df.createOrReplaceTempView("user")

    spark.udf.register("avgAge",functions.udaf(new MyAvgUDAF()))

    spark.sql("select avgAge(age) from user").show()
    /*A
    +--------------+
    |myavgudaf(age)|
    +--------------+
    |            30|
    +--------------+
     */
    // TODO 关闭环境
    spark.close()
  }
  /*
  自定义聚合函数类，计算平均年龄
  1.继承org.apache.spark.sql.expressions.Aggregator定义泛型
    IN:Long
    BUF:Buff
    OUT:Long
  2.重写方法
   */
  case class Buff(var total:Long,var count:Long)
  class MyAvgUDAF extends Aggregator[Long,Buff,Long] {
    //z & zero:初始值或零值
    //缓冲区初始化
    override def zero: Buff = {
      Buff(0L,0L)
    }

    //根据输入数据更新缓冲区数据
    override def reduce(b: Buff, a: Long): Buff = {
      b.total=b.total+a
      b.count=b.count+1
      b
    }

    //合并缓冲区
    override def merge(b1: Buff, b2: Buff): Buff = {
      b1.total=b1.total+b2.total
      b1.count=b1.count+b2.count
      b1
    }

    //计算结果
    override def finish(reduction: Buff): Long = {
      reduction.total/reduction.count
    }

    //缓冲区编码
    override def bufferEncoder: Encoder[Buff] = Encoders.product

    //输出编码
    override def outputEncoder: Encoder[Long] = Encoders.scalaLong
  }
}
